import timm
import torch.nn as nn
import torch
class ft_net_swin_extract(nn.Module):

    def __init__(self, class_num, droprate=0.5, stride=2,):
        super(ft_net_swin_extract, self).__init__()
        model_ft = timm.create_model('swin_base_patch4_window7_224', pretrained=True)
        # avg pooling to global pooling
        #model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.head = nn.Sequential() # save memory
        self.model = model_ft
    def forward(self, x):
        x = self.model.forward_features(x)
        return x